
import os
import sys
import time
import traceback
import multiprocessing as mp

import rpyc
from rpyc.utils import server

import dexen
from dexen.system import util, node
from dexen.system.base import service, MyManager, task
from dexen.system.base.spec import DexenSpec
#from dexen.system.base.exceptions import ScriptException, SlaveException
from dexen.system.node.slave.comm import SlaveComm

#===============================================================================
# Globals
#===============================================================================
MyManager.register("SlaveComm", SlaveComm)

#===============================================================================
# Starting
#===============================================================================
def start(_slavename, _nodename, _rreg_addr):
    global slavename, slavedir, nodename, manager, exec_proc, comm
    slavename = _slavename
    nodename = _nodename
    manager = MyManager()
    manager.start()
    
    slavedir = os.path.join(node.TEMP_DIR, slavename)
    if not os.path.exists(slavedir):
        os.mkdir(slavedir)
    sys.path.append(slavedir)
    
    comm = manager.SlaveComm(slavename, slavedir)

    print "The execute_tasks process is being initiated"
    exec_proc = mp.Process(target=execute_tasks, args=(comm,))
    exec_proc.start()

    SlaveService.ALIASES = [slavename]
    svr = server.ThreadedServer(SlaveService,
                                protocol_config=service.default_config)
    
    conn = rpyc.connect(_rreg_addr[0], _rreg_addr[1])
    my_addr = (service.get_my_ip_addr(), svr.port)
    conn.root.register_slave(slavename, my_addr)
    conn.close()
    
    svr.start()


            
#===============================================================================
# Service Class
#===============================================================================
class SlaveService(dexen.Service):
    def on_connect(self):
        pass
    
    def on_disconnect(self):
        pass

    def run(self, spec):
        #print "I am in slave these are the tasks", tasks
        spec = rpyc.classic.obtain(spec)
        assert isinstance (spec, DexenSpec) 
        self.mastername = spec.mastername
        self.masterdir = os.path.join(slavedir, self.mastername)
        spec.masterdir = self.masterdir
        if not os.path.exists(self.masterdir):
            os.mkdir(self.masterdir)
        #print "The master dir is", spec.masterdir
        util.unzip(spec.schema_zip, spec.masterdir)
        self.spec = spec
        comm.set_env(self.masterdir)

    def put(self, unit):
        sys.path.append(self.masterdir)
        cur = os.getcwd()
        os.chdir(self.masterdir)

        unit = rpyc.classic.obtain(unit)
        assert isinstance(unit, task.InputUnit)
        unit.mastername = self.mastername
        unit.masterdir = self.masterdir
        #print "in slave service thread, the masterdir", self.masterdir
        unit.obj = self.spec.get_task_obj(unit.task_name)
        #self.spec.objs = {} #empties the spec.objs
        #print "after getting the obj"
        #comm.set_env(self.masterdir)
        comm.put_input(unit)
        #comm.reset_env()
        sys.path = sys.path[:-1]
        os.chdir(cur)
        
    def get(self):
        return comm.get_output()
        
#===============================================================================
# TASK EXECUTOR
#===============================================================================
def execute_tasks(comm):
    
    while True:
        if not comm.has_input():
            time.sleep(1)
            continue

        masterdir = comm.get_env()
        sys.path.append(masterdir)
        os.chdir(masterdir)
        i_unit = comm.get_input()

        task_name = i_unit.task_name
        
        assert isinstance(i_unit, task.InputUnit) 
        print "executing", task_name
        
        #sys.path.append(i_unit.masterdir)
        #os.chdir(i_unit.masterdir)
        
        obj = i_unit.obj
        
        try:
            start_time = time.time()
            inds = obj.__execute__(i_unit.inds)
            exec_time = time.time() - start_time
        except:
            o_unit = task.OutputUnit(task_name)
            o_unit.crashed("crashed") 
        else:
            o_unit = task.OutputUnit(task_name, inds, exec_time)
        finally:
            comm.put_output(o_unit)    
            sys.path = sys.path[:-1]
#===============================================================================
# TESTS
#===============================================================================
def test_suite():
    pass

if __name__ == '__main__':
    test_suite()
    

